package main

import (
	"context"
	"errors"
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync/atomic"
	"syscall"
	"time"

	"golang.org/x/sync/semaphore"
)

var quit = make(chan os.Signal)

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var total int64 = 0
	pool, err := NewPool(ctx, 20, 10)
	if err != nil {
		log.Println(err)
	}

	for i := 0; i < 100; i++ {
		_ = pool.Put(func(ctx context.Context) {
			time.Sleep(time.Second * 3)
			atomic.AddInt64(&total, 1)
		})
	}

	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)

	go pool.Close()
	<-pool.closed
	<-quit

	fmt.Println(total)
}

type Pool struct {
	Ctx            context.Context
	taskCap        uint64
	taskCh         chan func(ctx context.Context)
	semaphore      *semaphore.Weighted
	state          uint32 //状态:0:停止,1:运行
	closed         chan int
	runningWorkers uint64 //正在运行的任务数
}

func NewPool(ctx context.Context, cap uint64, gNum uint64) (pool *Pool, err error) {
	if gNum > cap {
		return nil, errors.New("gNum cannot greater than cap")
	}

	pool = &Pool{
		Ctx:       ctx,
		taskCap:   cap,
		taskCh:    make(chan func(ctx context.Context), cap),
		semaphore: semaphore.NewWeighted(int64(gNum)),
		state:     1,
		closed:    make(chan int),
	}

	go func(ctx context.Context) {
		defer func() {
			if err := recover(); err != nil {
				log.Println(err)
			}
		}()
		for {
			select {
			case fn, ok := <-pool.taskCh:
				if !ok {
					return
				}
				pool.run(fn)
			case <-ctx.Done():
				return
			}
		}
	}(ctx)

	return
}

// 运行的任务数+1
func (p *Pool) incRunning() { // runningWorkers + 1
	atomic.AddUint64(&p.runningWorkers, 1)
}

// 运行的任务数-1
func (p *Pool) decRunning() { // runningWorkers - 1
	atomic.AddUint64(&p.runningWorkers, ^uint64(0))
}

func (pool *Pool) run(fn func(ctx context.Context)) {
	_ = pool.semaphore.Acquire(pool.Ctx, 1)
	defer pool.semaphore.Release(1)

	pool.incRunning()
	go func() {
		defer func() {
			pool.decRunning() //执行完任务后，任务数-1
			if err := recover(); err != nil {
				log.Println("fn exec error:(%+v)", err)
			}
		}()
		fn(pool.Ctx)
	}()
}

func (pool *Pool) Put(fn func(ctx context.Context)) (err error) {
	if !atomic.CompareAndSwapUint32(&pool.state, 1, 1) {
		return errors.New("task pool is already closed")
	}
	// 放入任务队列
	pool.taskCh <- fn
	return
}

func (pool *Pool) Close() {
	if !atomic.CompareAndSwapUint32(&pool.state, 1, 0) {
		return
	}
	// 等任务池中所有的任务执行完后再关闭任务池
	for len(pool.taskCh) > 0 || pool.runningWorkers > 0 {

	}
	pool.closed <- 1

	close(pool.taskCh)
	quit <- syscall.SIGINT
}
